package top.yudoge.dogeeditor.llm.tools

import dev.langchain4j.agent.tool.Tool
import dev.langchain4j.agent.tool.ToolMemoryId
import top.yudoge.dogeeditor.commons.cache.Cache
import top.yudoge.dogeeditor.commons.utils.convertList
import top.yudoge.dogeeditor.db.api.DataSourceMetaDataApi
import top.yudoge.dogeeditor.llm.constants.SessionMemoryKeyConstants.Companion.DB_NAME
import top.yudoge.dogeeditor.llm.constants.SessionMemoryKeyConstants.Companion.DS_NAME
import top.yudoge.dogeeditor.llm.constants.SessionMemoryKeyConstants.Companion.SCHEMA_NAME
import top.yudoge.dogeeditor.llm.memory.SessionMemory
import top.yudoge.dogeeditor.llm.objects.mro.TableMRO
import top.yudoge.dogeeditor.llm.task.TableEmbeddingTask
import top.yudoge.dogeeditor.llm.util.tokenizer.Tokenizer
import top.yudoge.dogeeditor.llm.util.VectorUtils

/**
 * Tools used by higher llm object like Agent which provide some database functionality.
 * Not specific for any llm object implementation.
 */
class DBTools(
    private val sessionMemory: SessionMemory,
    private val metaDataApi: DataSourceMetaDataApi,
    private val cache: Cache<String>, // global cache
    private val tokenizer: Tokenizer, // global tokenizer
) {

    @Tool("You can only use this tool if the `Selected table` is not empty! Pass the user selected table names in this tool directly. This tool will return the Table object.")
    fun getTableInfo(@ToolMemoryId sessionId: String, tableNames: List<String>): List<TableMRO> {

        val dsName = sessionMemory.get(sessionId, DS_NAME) ?: return emptyList()
        val dbName = sessionMemory.get(sessionId, DB_NAME) ?: return emptyList()
        val schemaName = sessionMemory.get(sessionId, SCHEMA_NAME) ?: return emptyList()

        return tableNames.map {
            TableMRO(
                it,
                metaDataApi.getColumns(dsName, dbName, schemaName, it).convertList()
            )
        }.filter { it.columns.isNotEmpty() }
    }

    @Tool("Given some words that appear in user task that maybe a part of some table name or remark. "+
            "it will match tables use these words. And return the Table objects that matched!")
    fun getPossibleTableInfo(@ToolMemoryId sessionId: String, possibleTableName: List<String>): List<TableMRO> {
        val possibleTableNames = getPossibleTableNames(sessionId, possibleTableName);
        return getTableInfo(sessionId, possibleTableNames)
    }

    private fun getPossibleTableNames(sessionId: String, possibleTableName: List<String>): List<String> {
//        val dsName = sessionMemory.get(sessionId, DS_NAME) ?: return emptyList()
//        val dbName = sessionMemory.get(sessionId, DB_NAME) ?: return emptyList()
//        val schemaName = sessionMemory.get(sessionId, SCHEMA_NAME) ?: return emptyList()
//
//        val ( vocabulary, allEmbeddings ) = tableEmbeddingService.getTableEmbeddings(dsName, dbName, schemaName)
//
//        return possibleTableName.map {
//            var cachedTableName = cache.get(cacheKey(dsName, dbName, schemaName, it), String::class.java)
//            if (cachedTableName != null) {
//                cachedTableName
//            } else {
//                val embedding = TableEmbeddingTask.embedding(tokenizer, it, vocabulary)
//                var bestSimilarIndex = VectorUtils.bestSimilar(embedding, allEmbeddings.map { it.embedding })
//                if (bestSimilarIndex != -1) {
//                    allEmbeddings[bestSimilarIndex].tableName
//                } else {
//                    ""
//                }
//            }
//        }.filter { it.isNotBlank() }
        return emptyList()
    }

    private fun cacheKey(dsName: String, dbName: String, schemaName: String, possibleTableName: String) =
        "dbtools:possible_table_name:$dsName:$dbName:$schemaName:$possibleTableName";

}